# remote_server.py
import sys
import socket
import threading
import struct
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QLabel, QPushButton,
    QVBoxLayout, QWidget, QHBoxLayout, QLineEdit, QMessageBox
)
from PyQt6.QtCore import Qt, QThread, pyqtSignal
from PyQt6.QtGui import QPixmap, QImage

# -------------------- 屏幕接收线程 --------------------
class ScreenReceiver(QThread):
    frame_received = pyqtSignal(QPixmap)
    connection_lost = pyqtSignal(str)

    def __init__(self, conn):
        super().__init__()
        self.conn = conn
        self.running = True

    def run(self):
        try:
            while self.running:
                len_data = self._recv_exact(4)
                if not len_data:
                    break
                data_len = struct.unpack("!I", len_data)[0]
                if data_len <= 0:
                    break

                img_data = self._recv_exact(data_len)
                if not img_data:
                    break

                image = QImage.fromData(img_data, "JPG")
                if image.isNull():
                    # 解析失败，跳过
                    continue
                pixmap = QPixmap.fromImage(image)
                self.frame_received.emit(pixmap)
        except Exception as e:
            self.connection_lost.emit(f"接收错误: {str(e)}")
        finally:
            self.connection_lost.emit("客户端断开连接")

    def _recv_exact(self, size):
        """确保按指定长度完整接收"""
        buf = b''
        while len(buf) < size:
            try:
                chunk = self.conn.recv(size - len(buf))
            except Exception:
                return None
            if not chunk:
                return None
            buf += chunk
        return buf

    def stop(self):
        # 不直接关闭 socket，让主线程做关闭操作
        self.running = False

# -------------------- 主窗口 --------------------
class RemoteControlServer(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("远程控制服务器 - PyQt6版")
        self.setMinimumSize(1024, 768)

        self.server_socket = None      # 图片通道监听 socket
        self.ctrl_socket = None        # 控制通道监听 socket
        self.client_conn = None        # 图片通道客户端连接
        self.ctrl_conn = None          # 控制通道客户端连接
        self.receiver_thread = None
        self.is_connected = False

        # 远端（客户端）实际屏幕分辨率（由客户端通过 INFO 发来）
        self.remote_width = 1280
        self.remote_height = 720

        # 发送锁（防止主线程和其它线程同时 send）
        self.send_lock = threading.Lock()

        self._build_ui()

    # -------------------- UI 构建 --------------------
    def _build_ui(self):
        main_layout = QVBoxLayout()

        top_bar = QHBoxLayout()
        self.ip_input = QLineEdit("0.0.0.0")
        self.port_input = QLineEdit("8888")
        self.start_btn = QPushButton("启动服务器")
        self.status_label = QLabel("状态：未启动")

        top_bar.addWidget(QLabel("服务器IP："))
        top_bar.addWidget(self.ip_input)
        top_bar.addWidget(QLabel("端口："))
        top_bar.addWidget(self.port_input)
        top_bar.addWidget(self.start_btn)
        top_bar.addWidget(self.status_label)

        self.screen_label = QLabel()
        self.screen_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.screen_label.setStyleSheet("background-color: black;")
        self.screen_label.setMinimumSize(800, 600)

        main_layout.addLayout(top_bar)
        main_layout.addWidget(self.screen_label, stretch=1)

        container = QWidget()
        container.setLayout(main_layout)
        self.setCentralWidget(container)

        self.start_btn.clicked.connect(self.toggle_server)
        # 鼠标事件
        self.screen_label.mouseMoveEvent = self.mouse_move
        self.screen_label.mousePressEvent = self.mouse_click

    # -------------------- 网络逻辑 --------------------
    def toggle_server(self):
        if self.is_connected or self.server_socket:
            self.stop_server()
        else:
            self.start_server()

    def start_server(self):
        try:
            ip = self.ip_input.text().strip()
            port = int(self.port_input.text().strip())
            # 图片通道监听
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.server_socket.bind((ip, port))
            self.server_socket.listen(1)

            # 控制通道监听（默认 port+1）
            ctrl_port = port + 1
            self.ctrl_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.ctrl_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.ctrl_socket.bind((ip, ctrl_port))
            self.ctrl_socket.listen(1)

            self.status_label.setText(f"等待客户端连接 {ip}:{port} (控制 {ctrl_port})")
            # 分别接受图片与控制通道（两个线程）
            threading.Thread(target=self.accept_image_client, daemon=True).start()
            threading.Thread(target=self.accept_ctrl_client, daemon=True).start()

            self.start_btn.setText("关闭服务器")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"启动失败: {e}")

    def accept_image_client(self):
        try:
            conn, addr = self.server_socket.accept()
            self.client_conn = conn
            self.is_connected = True
            self.status_label.setText(f"✅ 已连接 (图像): {addr}")
            self.receiver_thread = ScreenReceiver(conn)
            self.receiver_thread.frame_received.connect(self.update_screen)
            self.receiver_thread.connection_lost.connect(self.handle_disconnect)
            self.receiver_thread.start()
        except Exception as e:
            self.status_label.setText(f"图像连接错误: {e}")

    def accept_ctrl_client(self):
        try:
            # accept control connection (blocks until connected)
            conn, addr = self.ctrl_socket.accept()
            self.ctrl_conn = conn
            print("控制通道已连接:", addr)
            # 读取客户端可能发来的 INFO（客户端在连接后会发送 INFO,width,height）
            try:
                data = conn.recv(1024)
                if data:
                    s = data.decode(errors='ignore')
                    if s.startswith("INFO,"):
                        parts = s.strip().split(",")
                        if len(parts) >= 3:
                            try:
                                self.remote_width = int(parts[1])
                                self.remote_height = int(parts[2])
                                print("设置远端分辨率为:", self.remote_width, self.remote_height)
                            except:
                                pass
            except Exception:
                pass
        except Exception as e:
            print("控制通道 accept 错误:", e)

    def update_screen(self, pixmap):
        # 按保持比例缩放显示
        scaled = pixmap.scaled(
            self.screen_label.size(),
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation
        )
        self.screen_label.setPixmap(scaled)

    def handle_disconnect(self, msg):
        self.status_label.setText(f"⚠️ {msg}")
        self.stop_server()

    # -------------------- 鼠标控制 --------------------
    def mouse_move(self, event):
        if not self.ctrl_conn:
            return
        try:
            w, h = self.screen_label.width(), self.screen_label.height()
            if w == 0 or h == 0:
                return
            x_ratio = float(event.position().x()) / float(w)
            y_ratio = float(event.position().y()) / float(h)
            cmd = f"move,{x_ratio:.6f},{y_ratio:.6f}"
            with self.send_lock:
                self.ctrl_conn.sendall(cmd.encode())
        except Exception as e:
            print("发送鼠标移动失败:", e)

    def mouse_click(self, event):
        if not self.ctrl_conn:
            return
        try:
            w, h = self.screen_label.width(), self.screen_label.height()
            if w == 0 or h == 0:
                return
            x_ratio = float(event.position().x()) / float(w)
            y_ratio = float(event.position().y()) / float(h)
            cmd = f"click,{x_ratio:.6f},{y_ratio:.6f}"
            with self.send_lock:
                self.ctrl_conn.sendall(cmd.encode())
        except Exception as e:
            print("发送鼠标点击失败:", e)

    # -------------------- 停止与退出 --------------------
    def stop_server(self):
        self.is_connected = False
        self.start_btn.setText("启动服务器")

        if self.receiver_thread:
            self.receiver_thread.stop()
            self.receiver_thread.wait()
            self.receiver_thread = None

        if self.client_conn:
            try:
                self.client_conn.close()
            except:
                pass
            self.client_conn = None

        if self.ctrl_conn:
            try:
                self.ctrl_conn.close()
            except:
                pass
            self.ctrl_conn = None

        if self.server_socket:
            try:
                self.server_socket.close()
            except:
                pass
            self.server_socket = None

        if self.ctrl_socket:
            try:
                self.ctrl_socket.close()
            except:
                pass
            self.ctrl_socket = None

        self.status_label.setText("服务器已关闭")
        self.screen_label.clear()

    def closeEvent(self, event):
        self.stop_server()
        event.accept()

# -------------------- 程序入口 --------------------
if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = RemoteControlServer()
    window.show()
    sys.exit(app.exec())
